import os
import datetime
import tensorflow as tf
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense, Dropout, Input
from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard

# 图像大小
img_size = 100
img_channels = 1  # 灰度图
batch_size = 32
nb_classes = 2    # 两类：jump / nojump
epochs = 25

# 数据路径
data_dir = './dataset'

# 数据增强 + 归一化
datagen = ImageDataGenerator(
    rescale=1./255,
    validation_split=0.2,
    horizontal_flip=True,
    zoom_range=0.1,
)

# 训练集生成器
train_generator = datagen.flow_from_directory(
    data_dir,
    target_size=(img_size, img_size),
    batch_size=batch_size,
    color_mode='grayscale',
    class_mode='categorical',
    subset='training'
)

# 验证集生成器
validation_generator = datagen.flow_from_directory(
    data_dir,
    target_size=(img_size, img_size),
    batch_size=batch_size,
    color_mode='grayscale',
    class_mode='categorical',
    subset='validation'
)

print("类别映射:", train_generator.class_indices)

# 构建模型
model = Sequential()
# 使用 Input 层指定输入形状
model.add(Input(shape=(img_size, img_size, img_channels)))
model.add(Conv2D(32, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))

model.add(Conv2D(32, (3, 3), activation='relu'))
model.add(MaxPooling2D((2, 2)))

model.add(Flatten())
model.add(Dense(128, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(nb_classes, activation='softmax'))

# 编译模型
model.compile(optimizer='adam',
              loss='categorical_crossentropy',
              metrics=['accuracy'])

# 创建TensorBoard回调
log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
tensorboard_callback = TensorBoard(log_dir=log_dir, histogram_freq=1)

# 回调函数：保存最佳模型
checkpoint = ModelCheckpoint('newWeight.keras', save_best_only=True, monitor='val_accuracy', mode='max')

# 开始训练
history = model.fit(
    train_generator,
    epochs=epochs,
    validation_data=validation_generator,
    callbacks=[checkpoint, tensorboard_callback]  # 添加TensorBoard回调到callbacks列表中
)

# 可选：保存整个模型
model.save('dino_gesture_model.keras')

# 启动TensorBoard查看训练情况
# 在命令行运行:
# tensorboard --logdir logs/fit